package code;

import com.baomidou.mybatisplus.generator.FastAutoGenerator;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.builder.CustomFile;
import com.baomidou.mybatisplus.generator.config.rules.DbColumnType;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;

import java.nio.file.Paths;
import java.sql.Types;
import java.util.*;

public class CodeGenerator {

    public static void main(String[] args) {
        //模块名 (包名)
        String moduleName = "demo";
        //作者
        String author = "LX";
        //表名，多个英文逗号分隔
        String tableNames = "user";

        //------------------------------------------------------------

        String outputBaseDir = Paths.get(System.getProperty("user.dir")).toString();
        String modulePkg = "".equals(moduleName)? "" : "/" + moduleName;
        String moduleImportPkg = "".equals(moduleName)? "" : "." + moduleName;
        Map<OutputFile, String> pathInfoMap = new HashMap<>();
        pathInfoMap.put(OutputFile.entity, outputBaseDir+"/scaffold-myproject/src/main/java/com/scaffold/myproject/entity"+modulePkg);
        pathInfoMap.put(OutputFile.mapper, outputBaseDir+"/scaffold-myproject/src/main/java/com/scaffold/myproject/mapper"+modulePkg);
        pathInfoMap.put(OutputFile.service, outputBaseDir+"/scaffold-myproject/src/main/java/com/scaffold/myproject/service"+modulePkg);
        pathInfoMap.put(OutputFile.serviceImpl, outputBaseDir+"/scaffold-myproject/src/main/java/com/scaffold/myproject/service/impl"+modulePkg);
        pathInfoMap.put(OutputFile.xml, outputBaseDir+"/scaffold-myproject/src/main/resources/mapper"+modulePkg);
        FastAutoGenerator.create("jdbc:mysql://localhost:3306/testdb?serverTimezone=Asia/Shanghai&useUnicode=true&characterEncoding=UTF-8&useSSL=false",
                        "root",
                        "123456")
                .globalConfig((scanner,builder) -> {
                    builder.author(author) // 设置作者
                            .enableSwagger() // 开启 swagger 模式
                            .outputDir(outputBaseDir + "/scaffold-myproject/src/main/java/"); // 指定输出目录
                })
                .dataSourceConfig(builder ->
                        builder.typeConvertHandler((globalConfig, typeRegistry, metaInfo) -> {
                            int typeCode = metaInfo.getJdbcType().TYPE_CODE;
                            if (typeCode == Types.SMALLINT
                                    || typeCode == Types.TINYINT) {
                                // 自定义类型转换
                                return DbColumnType.INTEGER;
                            }
                            return typeRegistry.getColumnType(metaInfo);
                        })
                )
                .packageConfig((scanner,builder) -> builder
                        .parent("")
                        .entity("com.scaffold.myproject.entity"+moduleImportPkg)
                        .mapper("com.scaffold.myproject.mapper"+moduleImportPkg)
                        .service("com.scaffold.myproject.service"+moduleImportPkg)
                        .serviceImpl("com.scaffold.myproject.service.impl"+moduleImportPkg)
                        .pathInfo(pathInfoMap)
                )
                .strategyConfig((scanner, builder) ->
                        builder.addInclude(getTables(tableNames))
                                .entityBuilder()
                                .addSuperEntityColumns("id", "created_time", "modified_time","attributes","istest")
                                .formatFileName("%sEntity")
                                .javaTemplate("templates/entity.java")
                                .enableFileOverride()

                                .mapperBuilder()
                                .mapperTemplate("templates/mapper.java")
                                .mapperXmlTemplate("templates/mapper.xml")
                                .enableFileOverride()

                                .serviceBuilder()
                                .formatServiceFileName("%sService")
                                .formatServiceImplFileName("%sServiceImpl")
                                .serviceTemplate("templates/Service.java")
                                .serviceImplTemplate("templates/ServiceImpl.java")
                                .enableFileOverride()

                                .controllerBuilder()
                                .disable()

                                .build()
                )
                .injectionConfig(injectConfig -> {
                    Map<String,Object> customMap = new HashMap<>();
                    customMap.put("package.DTO","com.scaffold.myproject.model.dto"+moduleImportPkg);
                    customMap.put("package.VO","com.scaffold.myproject.model.vo"+moduleImportPkg);
                    injectConfig.customMap(customMap); //注入自定义属性
                    injectConfig.customFile(new CustomFile.Builder()
                            .fileName("VO.java") //文件名称
                            .formatNameFunction(tableInfo -> tableInfo.getEntityName().replaceAll("Entity",""))
                            .templatePath("templates/vo.java.ftl") //指定生成模板路径
                            .packageName("com.scaffold.myproject.model.vo"+moduleImportPkg)
                            .filePath(outputBaseDir+"/scaffold-myproject/src/main/java/")
                            .enableFileOverride()
                            .build());
                    injectConfig.customFile(new CustomFile.Builder()
                            .fileName("DTO.java") //文件名称
                            .formatNameFunction(tableInfo -> tableInfo.getEntityName().replaceAll("Entity",""))
                            .templatePath("templates/dto.java.ftl") //指定生成模板路径
                            .packageName("com.scaffold.myproject.model.dto"+moduleImportPkg)
                            .filePath(outputBaseDir+"/scaffold-myproject/src/main/java/")
                            .enableFileOverride()
                            .build());
                })
                .templateEngine(new FreemarkerTemplateEngine()) // 使用Freemarker引擎模板，默认的是Velocity引擎模板
                .execute();
    }

    // 处理 all 情况
    protected static List<String> getTables(String tables) {
        return "all".equals(tables) ? Collections.emptyList() : Arrays.asList(tables.split(","));
    }
}